using UnityEngine;
using UnityEditor;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
using Sirenix.OdinInspector;
using Sirenix.OdinInspector.Editor;
using Sirenix.Utilities.Editor;
[ExecuteAlways]
public class Gen_Irradiance_Map : MonoBehaviour
{
    //Irradiance_Map
    public ComputeShader shader;
    public ComputeShader Prefilshader;
    public ComputeShader LUTshader;
    //
    
    [Range(0,1)] [ShowInInspector]
    public static float IBL_Intensity;
    public bool Gen_Irr;
    public bool Gen_Spec;
    public bool Gen_Lut;
    
    public bool Gen_RenderTexture;
    public Cubemap envCube;
    // [InspectorButton("Bake")]
    //public bool bake;
    private static Cubemap irrand=null;
    private static Cubemap preSpec=null;
    private static Texture2D LUT=null;
    private static Texture2D DFGMultiScatteringLUT=null;
    private static Texture2D MultiScattering_LUT_Avg=null;
    
    [Button("生成")]
    private void Bake()
    {
        if(CheckOrCreat())
            Debug.Log("finish!");
    }
    [Button("设置")] 
    public static void SetIBL()
    {
        if (!irrand||!preSpec||!LUT)
        {
            irrand=AssetDatabase.LoadAssetAtPath<Cubemap>("Assets/HDRS/irrand.cubemap");
            preSpec=AssetDatabase.LoadAssetAtPath<Cubemap>("Assets/HDRS/prefilteredEnvCube.cubemap");
            LUT=AssetDatabase.LoadAssetAtPath<Texture2D>("Assets/HDRS/LUTs/LUT.png");
            DFGMultiScatteringLUT=AssetDatabase.LoadAssetAtPath<Texture2D>("Assets/HDRS/LUTs/LUT_MultiScatter.png");
            MultiScattering_LUT_Avg=AssetDatabase.LoadAssetAtPath<Texture2D>("Assets/HDRS/LUTs/Eavgtex.png");
            
        }
        
        
        
        Shader.SetGlobalTexture("IBL_irradianceMap",irrand);
        Shader.SetGlobalTexture("IBL_lobeSpecularMap",preSpec);
        Shader.SetGlobalTexture("IBL_BRDFMap",LUT);
        Shader.SetGlobalTexture("_DFGMultiScatteringLUT",DFGMultiScatteringLUT);
        Shader.SetGlobalTexture("MultiScattering_LUT_Avg",MultiScattering_LUT_Avg);
        Shader.SetGlobalFloat("IBL_Intensity",IBL_Intensity);
        
    }
    bool CheckOrCreat()
    {
        if (!LUTshader || !shader || !LUTshader || !envCube)
        {
            Debug.Log("Something is null! check a little");
            return false;
        }
        if(Gen_Irr)
            Bake_Irradiance_Map(this.envCube);
        if(Gen_Spec)
            PrefilterSpecularCubemap(envCube);
        if(Gen_Lut)
        {
            BakeLUTGGXMultiscatter();
            BakeLUTGGXAndCloth();
        }
        return true;
    }

    
    public static Cubemap GetCubeMapByCubeTexture2D(Color[][] recolor,int width)
    {
        
        int cubeMapSize = width;
        Cubemap cubemap = new Cubemap(cubeMapSize, TextureFormat.RGB24, false);
        for (int i = 0; i < 6; i++)
        {
            
            //cubemap.SetPixels(tempTex.GetPixels(cubeMapSize, 0, cubeMapSize, cubeMapSize), CubemapFace.PositiveY);
            cubemap.SetPixels(recolor[i], CubemapFace.PositiveX+i);
        }
        
       
        cubemap.Apply();
        return cubemap;
    }
    
    void Bake_Irradiance_Map(Cubemap cubemap)
    {
        int kernelHandle = shader.FindKernel("CSMain");

        
        int width = cubemap.width;
        if (width == 0)
        {
            Debug.Log("error width");
        }

        int size2 = width * width;
        Color[][] tempResult=new Color [6][];
        for (int i = 0; i < tempResult.Length; i++)
        { tempResult[i] = new Color[size2]; }
        ComputeBuffer [] resultBuffer=
        {
            new ComputeBuffer(size2,sizeof(float)*4),
            new ComputeBuffer(size2,sizeof(float)*4),
            new ComputeBuffer(size2,sizeof(float)*4),
            new ComputeBuffer(size2,sizeof(float)*4),
            new ComputeBuffer(size2,sizeof(float)*4),
            new ComputeBuffer(size2,sizeof(float)*4)
        };
        bool genRenderTex = this.Gen_RenderTexture;
        RenderTexture texarr=null;
        if (genRenderTex)
        {
            texarr= new RenderTexture(width, width,0,RenderTextureFormat.ARGB32);
            texarr.enableRandomWrite = true;
            texarr.dimension = TextureDimension.Tex2DArray;
            texarr.volumeDepth = 6;
            texarr.Create();
        }
        shader.SetBool("Gen_renderTexture",genRenderTex);

        if (Gen_RenderTexture)
        {
            shader.SetTexture(kernelHandle, "re_arr", texarr);
        }
        
        shader.SetTexture(kernelHandle, "cube_map", cubemap);
        for (int i = 0; i < 6; i++)
        {
            string s = "re_Result[" + i.ToString() + "]";
            shader.SetBuffer(kernelHandle,s,resultBuffer[i]);
        }
        
        shader.Dispatch(kernelHandle, width / 8, width / 8, 6);
        for (int i = 0; i < 6; i++)
        {
            resultBuffer[i].GetData(tempResult[i]);
        }

        var recubemap = GetCubeMapByCubeTexture2D(tempResult, width);
        for (int i = 0; i < 6; i++)
        {
            resultBuffer[i].Release();
        }

        if (Gen_RenderTexture)
        {
            AssetDatabase.CreateAsset(texarr, "Assets/HDRS/irrarr.renderTexture");

        }
        AssetDatabase.CreateAsset(recubemap, "Assets/HDRS/irrand.cubemap");
      
       
        AssetDatabase.Refresh();
    }
    void PrefilterSpecularCubemap(Cubemap cubemap)
    {
        int kernelHandle = Prefilshader.FindKernel("PrefilterEnv");
        int bakeSize = 128; 
        var environmentMap = new Cubemap(bakeSize, TextureFormat.RGBAFloat, true);
        int maxMip = environmentMap.mipmapCount;
        int sampleCubemapSize = cubemap.width;
        environmentMap.filterMode = FilterMode.Trilinear;
        for (int mip = 0; mip < maxMip; mip++)
        {
            int size = bakeSize >> mip;
            int size2 = size * size;
            Color[] tempResult = new Color[size2];
            float roughness = (float)mip / (float)(maxMip - 1);
            ComputeBuffer resultBuffer = new ComputeBuffer(size2, sizeof(float) * 4);
            for (int face = 0; face < 6; face++)
            {
                Prefilshader.SetInt("_Face", face);
                Prefilshader.SetTexture(kernelHandle, "_Cubemap", cubemap);
                Prefilshader.SetInt("_SampleCubemapSize", sampleCubemapSize);
                Prefilshader.SetInt("_Resolution", size);
                Debug.Log("roughness " + roughness);
                Prefilshader.SetFloat("_PerceptualRoughness", roughness);
                Prefilshader.SetBuffer(kernelHandle, "_Result", resultBuffer);
                Prefilshader.Dispatch(kernelHandle, size, size, 1);
                resultBuffer.GetData(tempResult);
                environmentMap.SetPixels(tempResult, (CubemapFace)face, mip);
            }
            resultBuffer.Release();
        }
        environmentMap.Apply(false);
        
        //Shader.SetGlobalTexture("_PrefilteredEnvMap", environmentMap);
        //
        AssetDatabase.CreateAsset(environmentMap, "Assets/HDRS/prefilteredEnvCube.cubemap");
        AssetDatabase.Refresh();
    }
  
    

    void BakeLUTGGXAndCloth()
    {
        int kernelHandle = LUTshader.FindKernel("DFV");
        RenderTexture renderTexture= new RenderTexture(512, 512,0,RenderTextureFormat.ARGB32);
        renderTexture.enableRandomWrite = true;
        renderTexture.dimension = TextureDimension.Tex2D;
        renderTexture.Create();
        {
            DFGMultiScatteringLUT=AssetDatabase.LoadAssetAtPath<Texture2D>("Assets/HDRS/LUTs/LUT_MultiScatter.png");
            MultiScattering_LUT_Avg=AssetDatabase.LoadAssetAtPath<Texture2D>("Assets/HDRS/LUTs/Eavgtex.png");
        }
        LUTshader.SetTexture(kernelHandle, "LUT", renderTexture);
        LUTshader.SetTexture(kernelHandle, "LUT_EuTex", DFGMultiScatteringLUT);
        LUTshader.SetTexture(kernelHandle, "LUT_EavgTex", MultiScattering_LUT_Avg);
        
        LUTshader.Dispatch(kernelHandle, 512/4, 512/4, 1);
        AssetDatabase.CreateAsset(renderTexture, "Assets/HDRS/LUTs/LUT.renderTexture");
        Debug.Log("LUT integration finished.");
        
        Texture2D texture2D = saveRenderTextureToTexture2D(renderTexture, TextureFormat.RGBA32);
        saveTexture2DToPNG(texture2D, "Assets/HDRS/LUTs/LUT.png");
        

    }
    void BakeLUTGGXMultiscatter()
    {
        int kernelHandle = LUTshader.FindKernel("DFV_Multiscatter");
        RenderTexture renderTexture = new RenderTexture(512, 512, 0,
            RenderTextureFormat.Default,RenderTextureReadWrite.Linear);
        renderTexture.enableRandomWrite = true;
        
        renderTexture.wrapMode = TextureWrapMode.Clamp;
        renderTexture.autoGenerateMips = false;
        renderTexture.useMipMap = false;
        renderTexture.Create();
        LUTshader.SetTexture(kernelHandle, "LUT_MultiScatter", renderTexture);
        LUTshader.Dispatch(kernelHandle, 512/4, 512/4, 1);
        
        kernelHandle = LUTshader.FindKernel("DFV_Eavg");
        RenderTexture Eavg = new RenderTexture(512, 512, 0,
            RenderTextureFormat.Default,RenderTextureReadWrite.Linear);
        Eavg.enableRandomWrite = true;
        Eavg.format = RenderTextureFormat.Default;
        Eavg.wrapMode = TextureWrapMode.Clamp;
        Eavg.autoGenerateMips = false;
        Eavg.useMipMap = false;
        Eavg.Create();
        
        LUTshader.SetTexture(kernelHandle, "LUT_Eu", renderTexture);
        LUTshader.SetTexture(kernelHandle, "LUT_Eavg", Eavg);
        LUTshader.Dispatch(kernelHandle, 1, 512/4, 1);

        
        Texture2D texture2D = saveRenderTextureToTexture2D(renderTexture, TextureFormat.RGBA32);
        saveTexture2DToPNG(texture2D, "Assets/HDRS/LUTs/LUT_MultiScatter.png");
        Texture2D Eavgtex = saveRenderTextureToTexture2D(Eavg, TextureFormat.RGBA32);
        saveTexture2DToPNG(Eavgtex, "Assets/HDRS/LUTs/Eavgtex.png");
    }
    public Texture2D saveRenderTextureToTexture2D(RenderTexture renderTexture, TextureFormat format)
    {
        RenderTexture.active = renderTexture;
        Texture2D texture2D = 
            new Texture2D(renderTexture.width, renderTexture.height, format, false, false);
        texture2D.wrapMode = TextureWrapMode.Clamp;
        texture2D.ReadPixels(new Rect(0, 0, renderTexture.width, renderTexture.height), 0, 0);
        texture2D.Apply();
        RenderTexture.active = null;
// #if UNITY_EDITOR
//         DestroyImmediate(renderTexture);
// #else
//             Destroy(renderTexture);
// #endif
        return texture2D;
    }
    
    public void saveTexture2DToPNG(Texture2D texture, string file) {
        byte[] bytes = texture.EncodeToPNG();
// #if UNITY_EDITOR
//         DestroyImmediate(texture);
// #else
//             Destroy(texture);
// #endif
        System.IO.File.WriteAllBytes(file, bytes);
        Debug.Log("Write tex2D to .png file.");
        AssetDatabase.Refresh();
    }


    
    

}
